#!/usr/bin/env python3
# Copyright (c) 2023-2024 maminjie <canpool@163.com>
# SPDX-License-Identifier: MulanPSL-2.0


import os
import argparse
import logging
import datetime
import platform
import shutil

logging.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=logging.INFO)

CUR_DIR = os.path.dirname(os.path.realpath(__file__))
YUM_CONFIG = ""

yum_config_template = """\
[main]
gpgcheck=1
installonly_limit=3
clean_requirements_on_remove=True
best=True
skip_if_unavailable=False
reposdir={reposdir}
"""

yum_repo_template = """\
[OS]
name=OS
baseurl=http://repo.openeuler.org/{oeversion}/OS/$basearch/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/OS/$basearch/RPM-GPG-KEY-openEuler

[everything]
name=everything
baseurl=http://repo.openeuler.org/{oeversion}/everything/$basearch/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/everything/$basearch/RPM-GPG-KEY-openEuler

[EPOL]
name=EPOL
baseurl=http://repo.openeuler.org/{oeversion}/EPOL/main/$basearch/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/OS/$basearch/RPM-GPG-KEY-openEuler

[EPOL-UPDATE]
name=EPOL-UPDATE
baseurl=http://repo.openeuler.org/{oeversion}/EPOL/update/main/$basearch/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/OS/$basearch/RPM-GPG-KEY-openEuler

[EPOL-UPDATE-source]
name=EPOL-UPDATE-source
baseurl=http://repo.openeuler.org/{oeversion}/EPOL/update/main/source/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/OS/$basearch/RPM-GPG-KEY-openEuler

[debuginfo]
name=debuginfo
baseurl=http://repo.openeuler.org/{oeversion}/debuginfo/$basearch/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/debuginfo/$basearch/RPM-GPG-KEY-openEuler

[source]
name=source
baseurl=http://repo.openeuler.org/{oeversion}/source/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/source/RPM-GPG-KEY-openEuler

[update]
name=update
baseurl=http://repo.openeuler.org/{oeversion}/update/$basearch/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/OS/$basearch/RPM-GPG-KEY-openEuler

[update-source]
name=update-source
baseurl=http://repo.openeuler.org/{oeversion}/update/source/
enabled=1
gpgcheck=1
gpgkey=http://repo.openeuler.org/{oeversion}/OS/$basearch/RPM-GPG-KEY-openEuler

"""


def create_repo_tree(dir):
    try:
        for d in ("aarch64", "x86_64", "source"):
            os.makedirs(os.path.join(dir, d, "Packages"), exist_ok=True)
    except OSError:
        return False
    return True


def create_oe_yum_repo(tmp_dir, version):
    global YUM_CONFIG
    yum_repo_dir = os.path.join(tmp_dir, "yum.repos.d")
    os.makedirs(yum_repo_dir, exist_ok=True)
    YUM_CONFIG = os.path.join(tmp_dir, "yum.conf")
    yum_repo_file = os.path.join(yum_repo_dir, "os.repo")
    with open(YUM_CONFIG, "w", encoding="utf-8") as f:
        f.write(yum_config_template.format(reposdir=yum_repo_dir))
    with open(yum_repo_file, "w", encoding="utf-8") as f:
        f.write(yum_repo_template.format(oeversion=version))


def get_pkg_list(pkg_file):
    pkgs = []
    if not os.path.isfile(pkg_file):
        return pkgs
    with open(pkg_file, 'r', encoding="utf-8") as f:
        pkgs = [x.strip() for x in f.readlines() if x.strip()]
    return pkgs


def do_picker(args):
    if args.dir:
        pick_dir = os.path.realpath(args.dir)
        if not os.path.isdir(pick_dir):
            logging.error("{} is not directory".format(args.dir))
            return False
    else:
        dt = datetime.datetime.now().strftime("%Y%m%d")
        pick_dir = "pick_{}".format(dt)
        pick_dir = os.path.join(CUR_DIR, pick_dir)
    if not create_repo_tree(pick_dir):
        logging.error("create pick tree failed")
        return False
    pkgs = []
    if args.pkg:
        pkgs = args.pkg
    elif args.file:
        pkgs = get_pkg_list(os.path.realpath(args.file))
    if not pkgs:
        logging.error("no packages")
        return False
    tmp_dir = os.path.join(pick_dir, "tmp")
    create_oe_yum_repo(tmp_dir, args.version)
    """x86_64 platform can't resolve the provider of aarch64 package
        yumdownloader --url fence-agents --repo openEuler_20.03_LTS_SP1_Epol_aarch64
    can't resolve the provider fence-agents-all
    """
    arch = platform.uname().machine
    if arch not in ("x86_64", "aarch64"):
        logging.error("arch {} is not supported now".format(arch))
        return False
    ret = True
    for arch in (arch, ):
        rpm_bin_file = os.path.join(tmp_dir, "rpm_bin_{}.txt".format(arch))
        rpm_src_file = os.path.join(tmp_dir, "rpm_src_{}.txt".format(arch))
        # get the package list
        cmd = "dnf download -c {} --resolve --url {}".format(YUM_CONFIG, " ".join(pkgs))
        cmd += " | grep '/update/'"
        cmd += " | awk -F '/' '{print $NF}' | sed -r 's/-[^-]+-[^-]+$//g'"
        cmd += " | tee {}".format(rpm_bin_file)
        if os.system(cmd):
            logging.error("{} failed".format(cmd))
            ret = False
            continue
        # download bin packages
        pick_pkgs = []
        with open(rpm_bin_file, "r", encoding="utf-8") as f:
            pick_pkgs = [x.strip() for x in f.readlines() if x.strip()]
        if not pick_pkgs:
            logging.error("dnf parse package from {} repo failed".format(arch))
            ret = False
            continue
        cmd = "dnf download -c {} {}".format(YUM_CONFIG, " ".join(pick_pkgs))
        cmd += " --destdir={}".format(os.path.join(pick_dir, arch, "Packages"))
        os.system(cmd)
        cmd += " --debuginfo --debugsource"
        os.system(cmd)
        # get source package list
        cmd = "for p in $(ls {}/{}/Packages/*.rpm | grep -v src.rpm); do rpm -qpi $p | grep '^Source RPM' | " \
              "cut -d ' ' -f 5 | sed -r 's/-[^-]+-[^-]+$//g'; done | sort -u".format(pick_dir, arch)
        cmd += " > {}".format(rpm_src_file)
        if os.system(cmd):
            logging.error("{} failed".format(cmd))
            ret = False
            continue
        # download source packages
        pick_src_pkgs = []
        with open(rpm_src_file, "r", encoding="utf-8") as f:
            pick_src_pkgs = [x.strip() for x in f.readlines() if x.strip()]
        if not pick_src_pkgs:
            logging.error("dnf parse package from {} repo failed".format(arch))
            ret = False
            continue
        cmd = "dnf download -c {} {}".format(YUM_CONFIG, " ".join(pick_src_pkgs))
        cmd += " --destdir={}".format(os.path.join(pick_dir, "source", "Packages"))
        cmd += " --source"
        os.system(cmd)
    if ret:
        shutil.rmtree(tmp_dir)
    return ret


def create_repo(dir):
    if os.path.isdir(dir):
        os.system("createrepo {}".format(dir))


def _publish_repo(dir):
    for d in ("aarch64", "x86_64", "source"):
        create_repo(os.path.join(dir, d))


def do_publish(args):
    repo_dir = os.path.realpath(args.dir)
    if not os.path.isdir(repo_dir):
        logging.error("{}: no such directory".format(args.dir))
        return False
    _publish_repo(repo_dir)
    return True


def do_main(args):
    print("try -h/--help for more details.")


def main():
    parser = argparse.ArgumentParser(description="create and publish repo")
    parser.set_defaults(func=do_main)

    subparsers = parser.add_subparsers(help="sub-commands")
    # picker
    subparser = subparsers.add_parser("picker", aliases=["pick"], formatter_class=argparse.RawTextHelpFormatter,
                                      help="pick packages and their dependencies")
    subparser.add_argument("version", type=str, metavar="VERSION", help="openEuler version, such as: openEuler-23.03")
    group = subparser.add_mutually_exclusive_group()
    group.add_argument("-p", "--pkg", type=str, nargs='+', metavar="PACKAGE", help="package name, support multiple")
    group.add_argument("-f", "--file", type=str, metavar="FILE", help="packages file (one per line)")
    subparser.add_argument("--dir", type=str, metavar="DIR", help="pick directory, default is pick_yyyymmdd")
    subparser.set_defaults(func=do_picker)
    # publish
    subparser = subparsers.add_parser("publish", aliases=["pub"], formatter_class=argparse.RawTextHelpFormatter,
                                      help="create repo data")
    subparser.add_argument("dir", type=str, metavar="DIR", help="repo parent directory")
    subparser.set_defaults(func=do_publish)

    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()
